import torch
import torch.nn as nn
import torch.nn.functional as F

class channel_mask_generator(nn.Module):
    """
    根据输入BoW向量和词嵌入，生成一个基于语义的0-1掩码。
    它计算每个词与文档整体语义的余弦相似度，并根据阈值进行筛选。
    """
    def __init__(self, word_embeddings, threshold=0.4):
        """
        word_embeddings: 模型的词嵌入矩阵 [vocab_size, embed_dim]
        threshold: 余弦相似度的阈值
        """
        super(channel_mask_generator, self).__init__()
        # 我们需要非参数化的词嵌入来进行计算
        self.word_embeddings = nn.Parameter(word_embeddings, requires_grad=False)
        self.threshold = threshold

    def forward(self, bow):
        """
        bow: 输入的词袋向量 [batch_size, vocab_size]
        """
        # 1. 计算文档的语义嵌入
        # 通过对词嵌入进行加权平均得到文档嵌入，权重为BoW中的词频
        # [batch_size, vocab_size] -> [batch_size, vocab_size, 1]
        bow_weights = F.normalize(bow, p=1, dim=1).unsqueeze(2)
        # [batch_size, vocab_size, 1] * [vocab_size, embed_dim] -> [batch_size, embed_dim]
        doc_embed = torch.sum(bow_weights * self.word_embeddings, dim=1)

        # 2. 计算文档嵌入与所有词嵌入的余弦相似度
        # 归一化
        doc_embed_norm = F.normalize(doc_embed, p=2, dim=1)
        word_embeddings_norm = F.normalize(self.word_embeddings, p=2, dim=0)
        # [batch_size, embed_dim] @ [embed_dim, vocab_size] -> [batch_size, vocab_size]
        cosine_sim = torch.matmul(doc_embed_norm, word_embeddings_norm.T)

        # 3. 根据阈值生成掩码
        # 大于阈值的为1，小于等于的为0
        semantic_mask = (cosine_sim > self.threshold).float()

        # 4. 确保只在原始文档存在的词上应用掩码
        final_mask = semantic_mask * (bow > 0).float()

        return final_mask
